import numpy as np
from scipy.spatial.distance import cdist

# 1. 模拟数据场景
def generate_scenario_data(scenario="low_dim"):
    """ 
根据场景生成数据
:param scenario: 场景类型，支持 'low_dim' 或 'high_dim'
:return: 数据库和查询向量
""" 
    np.random.seed(42)
    if scenario == "low_dim":
        database = np.random.rand(10, 3)    # 低维数据，3维
        query = np.random.rand(1, 3)    # 查询向量
    elif scenario == "high_dim":
        database = np.random.rand(10, 100)    # 高维稀疏数据，100维
        query = np.random.rand(1, 100)    # 查询向量
    else:
        raise ValueError("不支持的场景类型")
    return database, query

# 生成低维场景数据
low_dim_db, low_dim_query = generate_scenario_data("low_dim")

# 生成高维场景数据
high_dim_db, high_dim_query = generate_scenario_data("high_dim")

# 2. 欧氏距离与余弦相似度的计算
def calculate_metrics(database, query):
    """ 
计算欧氏距离和余弦相似度
:param database: 数据库向量
:param query: 查询向量
:return: 欧氏距离和余弦相似度
""" 
    euclidean_distances = cdist(query, database, metric="euclidean").flatten()
    cosine_similarities = 1 - cdist(query, database, metric="cosine").flatten()
    return euclidean_distances, cosine_similarities

low_dim_euclidean, low_dim_cosine = calculate_metrics(low_dim_db, low_dim_query)
high_dim_euclidean, high_dim_cosine = calculate_metrics(high_dim_db, high_dim_query)

# 3. 输出结果
print("低维场景（3维）:")
print("数据库向量:\n", low_dim_db)
print("查询向量:\n", low_dim_query.flatten())
print("欧氏距离:\n", low_dim_euclidean)
print("余弦相似度:\n", low_dim_cosine)

print("\n高维场景（100维）:")
print("数据库向量（前3个）:\n", high_dim_db[:3]) # 输出前 3 个向量以节省空间
print("查询向量（前10个值）:\n", high_dim_query.flatten()[:10], "...") # 输出前 10 个值
print("欧氏距离（前3个）:\n", high_dim_euclidean[:3], "...")
print("余弦相似度（前3个）:\n", high_dim_cosine[:3], "...")